package miner;

import org.json.JSONArray;
import org.json.JSONObject;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadPoolExecutor;

public class StratumMiner {

    //public-pool.io is the only solo pool I've discovered that allows me to customize the difficulty
    //Our Java miner has a low hash rate, if the difficulty of pools is too high, We cannot submit effective work for a long time and may be disconnected
    private static final String POOL_HOST = "public-pool.io";
    private static final int POOL_PORT = 21496;
//    private static final String POOL_HOST = "solo.ckpool.org";
//    private static final String POOL_HOST = "pool4ever.com";
//    private static final String POOL_HOST = "minorpool.com";
//    private static final String POOL_HOST = "us.emcd.network";
//    private static final String POOL_HOST = "us.vkbit.com";
//    private static final String POOL_HOST = "stratum.braiins.com";
//    private static final String POOL_HOST = "btc.bullpool.io";
//    private static final int POOL_PORT = 3333;

    //set to your btc address.workerName
    private static final String USERNAME = "bc1qst559j7hf7xfasnlvs8sdjuvltss0at0h0h7pk.workerLcj01";
    private static final String PASSWORD = "123456";

    public static String Extranonce1;
    public static int Extranonce2_size;

    public static BufferedReader in;
    public static OutputStreamWriter out;

    public static long id = 4;

    //Difficulty of sending suggestions to the mining pool. Not all mining pools support custom difficulty
    public static double startDiff = 0.08;

    //Current mining pool difficulty
    public static double difficulty = 1.0;
    public static BigInteger targetBig = new BigInteger("00000000FFFF0000000000000000000000000000000000000000000000000000", 16);

    //Set to your CPU core count
    public static final int MINER_THREAD_SIZE = 10;
    public static long[] hashCountList = new long[MINER_THREAD_SIZE];

    public static long startTime = 0;

    public static boolean isMinning = false;
    public static byte[] targetBytes;

    private static Socket socket;

    private static ExecutorService executorService;

    public static void main(String[] args) throws IOException {
        Arrays.fill(hashCountList, 0);
        executorService = Executors.newFixedThreadPool(MINER_THREAD_SIZE);
        reConnect();
        stats();
    }

    public static void reConnect(){
        isMinning = false;
        startTime = System.currentTimeMillis();
        while (true) {
            try {
                connectToPool();
                subscribe();
                authorize();
                // If we reach here, the connection was successful, so we break the loop
                break;
            } catch (IOException e) {
                System.err.println("Connection failed, retrying...");
                try {
                    Thread.sleep(2000); // Wait 2 seconds before retrying
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                    System.err.println("Thread interrupted, exiting...");
                    return;
                }
            }
        }
        listen();
        sendDifficulty(startDiff);
    }

    private static void connectToPool() throws IOException {
        socket = new Socket(POOL_HOST, POOL_PORT);
        in = new BufferedReader(new InputStreamReader(socket.getInputStream(), StandardCharsets.UTF_8));
        out = new OutputStreamWriter(socket.getOutputStream(), StandardCharsets.UTF_8);
    }

    public static void subscribe() throws IOException {
        JSONObject subscribeRequest = new JSONObject();
        subscribeRequest.put("id", 1);
        subscribeRequest.put("method", "mining.subscribe");
        subscribeRequest.put("params", new Object[]{"XiongErMiner"});

        out.write(subscribeRequest.toString() + "\n");
        out.flush();

        String subscribeResponse = in.readLine();
        System.out.println("Subscribe Response: " + subscribeResponse);

        if (subscribeResponse != null) {
            System.out.println("Subscribe success");
            JSONObject message = new JSONObject(subscribeResponse);
            if (message.has("result")) {
                Extranonce1 = message.getJSONArray("result").getString(1);
                Extranonce2_size = message.getJSONArray("result").getInt(2);
                System.out.println("Extranonce1: " + Extranonce1);
                System.out.println("Extranonce2_size: " + Extranonce2_size);
            }

        } else {
            System.out.println("Subscribe failed, retrying...");
            throw new IOException("Server returned null response");
        }
    }

    public static void authorize() throws IOException {
        JSONObject authorizeRequest = new JSONObject();
        authorizeRequest.put("id", 2);
        authorizeRequest.put("method", "mining.authorize");
        authorizeRequest.put("params", new Object[]{USERNAME, PASSWORD});

        out.write(authorizeRequest.toString() + "\n");
        out.flush();
        String authorizeResponse = in.readLine();

        System.out.println("Authorize Response: " + authorizeResponse);

        if (authorizeResponse != null) {
            System.out.println("Authorize success");
            JSONObject message = new JSONObject(authorizeResponse);
            if (message.has("result")) {

            }else if (message.has("method")) {

                String method = message.getString("method");
                if(method.equals("mining.set_difficulty")){
                    if (message.has("params")){
                        double diff = message.getJSONArray("params").getDouble(0);
                        setDifficulty(diff);
                    }
                }

            }
        }else {
            System.out.println("Authorize failed, retrying...");
            throw new IOException("Server returned null response");
        }
    }

    public static void sendDifficulty(double diff){

        JSONObject diffRequest = new JSONObject();
        diffRequest.put("id", 114514);
        diffRequest.put("method", "mining.suggest_difficulty");
        diffRequest.put("params", new Object[]{diff});
        try {
            out.write(diffRequest.toString() + "\n");
            out.flush();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        System.out.println("submit: " + diffRequest.toString());
    }

    public static void listen() {
        Timer timer = new Timer();
        TimerTask task = new TimerTask() {
            @Override
            public void run() {
                try {
                    String line = in.readLine();
                    System.out.println("Receive Message: " + line);
                    if (line != null) {
                        handleMessage(line);
                        listen();
                    }else {
                        System.out.println("Receive Message failed, retrying...");
                        reConnect();
                    }
                } catch (IOException e) {

                }

            }
        };
        timer.schedule(task, 0);
    }

    public static void stats(){
        Timer timer = new Timer();
        TimerTask task = new TimerTask() {
            @Override
            public void run() {
                try {
                    long count = 0;
                    for (int i = 0; i < hashCountList.length; i++) {
                        count += hashCountList[i];
                    }
                    double hashRate = count/((System.currentTimeMillis()-startTime)/1000);
                    System.out.println("hashRate: " + hashRate/1000000 + "MH/s");

                    //Count active threads
                    ThreadPoolExecutor threadPoolExecutor = (ThreadPoolExecutor) executorService;
                    int activeCount = threadPoolExecutor.getActiveCount();
                    System.out.println("Number of active threads: " + activeCount);
                }catch (Exception e){

                }
                stats();
            }
        };
        timer.schedule(task, 5000);
    }

    public static void handleMessage(String line){

        JSONObject message = new JSONObject(line);
        if (message.has("method")) {

            String method = message.getString("method");
            if(method.equals("mining.set_difficulty")){
                if (message.has("params")){
                    double diff = message.getJSONArray("params").getDouble(0);
                    setDifficulty(diff);
                }
            }else if(method.equals("mining.notify")){
                if (message.has("params")){
                    handleMiningNotify(message);
                }
            }

        }

    }

    public static void setDifficulty(double diff) {
        difficulty = diff;
        System.out.println("difficulty set to " + difficulty);

        String targetInitialHex = "00000000FFFF0000000000000000000000000000000000000000000000000000";
        BigInteger targetInitial = new BigInteger(targetInitialHex, 16);
        BigDecimal difficultyDecimal = new BigDecimal(difficulty);
        BigDecimal targetDecimal = new BigDecimal(targetInitial).divide(difficultyDecimal, 0, BigDecimal.ROUND_DOWN);
        BigInteger target = targetDecimal.toBigInteger();
        String targetHex = target.toString(16);
        while (targetHex.length() < 64) {
            targetHex = "0" + targetHex;
        }
        targetBig = new BigInteger(targetHex, 16);
        targetBytes = hexStringToByteArray(targetHex);
        System.out.println("target set to " + targetHex);
    }

    public static String generateHexString(int size) {

        byte[] bytes = new byte[size];
        new java.util.Random().nextBytes(bytes);
        return bytesToHex(bytes);
    }

    private static void handleMiningNotify(JSONObject message) {
        try {
            String jobId = message.getJSONArray("params").getString(0);
            String prevHash = message.getJSONArray("params").getString(1);
            String coinb1 = message.getJSONArray("params").getString(2);
            String coinb2 = message.getJSONArray("params").getString(3);
            JSONArray merkleBranch = message.getJSONArray("params").getJSONArray(4);
            String version = message.getJSONArray("params").getString(5);
            String nBits = message.getJSONArray("params").getString(6);
            String nTime = message.getJSONArray("params").getString(7);
            boolean cleanJobs = message.getJSONArray("params").getBoolean(8);

            if(!cleanJobs && isMinning){
                return;
            }
            isMinning = true;

            executorService.shutdownNow();
            executorService = Executors.newFixedThreadPool(MINER_THREAD_SIZE);

            String prevHashSmall = toSmallEndFor64Hex(prevHash);
            String versionSmall = toSmallEnd(version);

            String[] merkle_branch = new String[merkleBranch.length()];
            for (int i = 0; i < merkleBranch.length(); i++) {
                merkle_branch[i] = merkleBranch.getString(i);
            }
            String ntimeSmall = toSmallEnd(nTime);
            String nbitsSmall = toSmallEnd(nBits);

            for (int i = 0; i < MINER_THREAD_SIZE; i++) {
                int finalI = i;
                executorService.submit(() -> {
                    try {
                        miner(versionSmall, prevHashSmall, ntimeSmall, nbitsSmall, jobId, nTime, coinb1, coinb2, merkle_branch, finalI);
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                });
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void miner(String versionSmall, String prevHashSmall, String ntimeSmall, String nbitsSmall, String jobId, String nTime, String coinb1, String coinb2, String[] merkle_branch, int threadIdx) throws Exception {

        System.out.println("start Job: " + jobId);

        while (true){
            if (Thread.currentThread().isInterrupted()) {
                System.out.println("end Job: " + jobId);
                return;
            }

            String extranonce2 = generateHexString(Extranonce2_size);
            System.out.println("extranonce2: " + extranonce2);

            String coinbase = coinb1 + Extranonce1 + extranonce2 + coinb2;
            String merkleRoot = getMerkle_root00(coinbase, merkle_branch);

            String headPrefix = versionSmall + prevHashSmall + merkleRoot + ntimeSmall + nbitsSmall;
            byte[] headPrefixBytes = hexStringToByteArray(headPrefix);

            int nonce = 0;
            byte[] headBytes = new byte[headPrefixBytes.length+4];
            System.arraycopy(headPrefixBytes,0,headBytes,0,headPrefixBytes.length);


            while (true){
                if (Thread.currentThread().isInterrupted()) {
                    System.out.println("end Job: " + jobId);
                    return;
                }

                byte[] nonceBytes = int2bytes(nonce);
                System.arraycopy(nonceBytes,0,headBytes,headPrefixBytes.length,nonceBytes.length);
                byte[] hash = sha256D(headBytes);

                hashCountList[threadIdx] ++;
                if(checkHash(hash)){

                    String headHash = bytesToHex(hash);
                    String headHashBigEnd = toSmallEnd(headHash);
                    System.out.println(headHashBigEnd);
                    JSONObject submitRequest = new JSONObject();
                    submitRequest.put("id", id++);
                    submitRequest.put("method", "mining.submit");
                    submitRequest.put("params", new Object[]{
                            USERNAME,
                            jobId,
                            extranonce2,
                            nTime,
                            toSmallEnd(bytesToHex(nonceBytes))
                    });

                    out.write(submitRequest.toString() + "\n");
                    out.flush();

                    System.out.println("submit: " + submitRequest.toString());

                }

                nonce++;
                if(nonce>2147483646){
                    System.out.println("nonce use out");
                    break;
                }
            }

        }


    }

    public static byte[] int2bytes(int value){
        byte[] byteArray = new byte[4];
        byteArray[0] = (byte) (value >> 24);
        byteArray[1] = (byte) (value >> 16);
        byteArray[2] = (byte) (value >> 8);
        byteArray[3] = (byte) value;

        return byteArray;
    }

    public static boolean checkHash(byte [] hash){
        for (int i = 0; i < targetBytes.length; i++) {
            int unsignedT = targetBytes[i] & 0xFF;
            int unsignedH = hash[targetBytes.length-1-i] & 0xFF;
            if(unsignedH<unsignedT){
                return true;
            }else if(unsignedH>unsignedT){
                return false;
            }
        }
        return true;
    }



    public static String getMerkle_root00(String coinbase, String[] merkle_branch) throws Exception {

        byte [] res = sha256D(hexStringToByteArray(coinbase));

        for(String s : merkle_branch){
            byte[] sByte = hexStringToByteArray(s);
            byte [] temp = new byte[64];
            System.arraycopy(res,0,temp,0,32);
            System.arraycopy(sByte,0,temp,32,32);
            res = sha256D(temp);
        }

        String bigEnd = bytesToHex(res);
        return bigEnd;
    }

    private static byte[] hexStringToByteArray(String s) {
        int len = s.length();
        byte[] data = new byte[len / 2];
        for (int i = 0; i < len; i += 2) {
            data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4)
                    + Character.digit(s.charAt(i + 1), 16));
        }
        return data;
    }

    private static String bytesToHex(byte[] bytes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            sb.append(String.format("%02x", b));
        }
        return sb.toString();
    }

    private static byte[] sha256D(byte[] input){
        MessageDigest digest;
        try {
            digest = MessageDigest.getInstance("SHA-256");
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
        byte[] res1 = digest.digest(input);
        byte[] res2 = digest.digest(res1);
        return res2;
    }

    public static String toSmallEndFor64Hex(String input) {
        if(input.length()%8!=0){
            return "";
        }
        StringBuilder output = new StringBuilder();
        for(int i = input.length()-8; i >= 0; i-=8) {
            for (int j = 0; j < 8; j++) {
                output.append(input.charAt(i+j));
            }
        }
        String bigEnd = output.toString();
        String smallEnd = toSmallEnd(bigEnd);
        return smallEnd;
    }

    public static String toSmallEnd(String input) {

        if(input.length()%2!=0){
            return "";
        }
        StringBuilder output = new StringBuilder();
        for(int i = input.length()-2; i >= 0; i-=2) {
            output.append(input.charAt(i));
            output.append(input.charAt(i+1));
        }
        return output.toString();


    }
}
